#include <iostream>
#include <winsock2.h>
#include <windows.h>
#include <ws2tcpip.h>
#include <vector>
#include <string>
#include <algorithm>
#include <limits>
#include <sstream>
#include <shlobj.h>

// -------------------------------------------------------------------------
// 1. IOCTL 和设备常量
// -------------------------------------------------------------------------

#define DEVICE_NAME L"\\\\.\\tdnetfilter"

// 根据 FUN_140001024 逆向分析得到的 IOCTL 码
#define IOCTL_SET_WHITELIST (DWORD)0x120008 // 对应内核 Mode 1 (排除)
#define IOCTL_SET_BLACKLIST (DWORD)0x12000c // 对应内核 Mode 2 (包含)
#define IOCTL_SET_RESET     (DWORD)0x120018 // 对应内核 Mode 0 (软关闭/重置)

// -------------------------------------------------------------------------
// 2. 配置结构体 (0x98 字节)
// -------------------------------------------------------------------------

#define MAX_IP_COUNT 49      // 0x31 个 DWORD
#define MAX_LIST_NAME_SIZE 0x94 // 列表描述大小

// 匹配内核期望的内存块结构 (0x98 字节 = 0x94 + 0x4 + 0x31*0x4)
typedef struct {
	char szListName[MAX_LIST_NAME_SIZE]; // 列表描述 (0x0 - 0x93)
	DWORD dwIpCount;                     // IP 数量 (0x94)
	DWORD IpArray[MAX_IP_COUNT];         // IP 地址数组 (0x98 开始)
} CONFIG_PAYLOAD;

// -------------------------------------------------------------------------
// 3. 核心功能类
// -------------------------------------------------------------------------

class TDNetFilterController {
private:
	HANDLE hDevice = INVALID_HANDLE_VALUE;

	bool OpenDriverHandle() {
		if (hDevice != INVALID_HANDLE_VALUE) return true;

		hDevice = CreateFileW (
		              DEVICE_NAME,
		              GENERIC_READ | GENERIC_WRITE, // 0xC0000000
		              FILE_SHARE_READ | FILE_SHARE_WRITE, // 3
		              nullptr, // LPSECURITY_ATTRIBUTES
		              OPEN_EXISTING, // 3
		              0,
		              nullptr
		          );

		if (hDevice == INVALID_HANDLE_VALUE) {
			DWORD error = GetLastError();
			std::cerr << "[-] 无法打开驱动程序句柄 (" << DEVICE_NAME << ")。错误码: " << error << std::endl;

			// 提示权限问题
			if (error == 5) { // ERROR_ACCESS_DENIED
				std::cerr << "    请确认程序是否以管理员权限运行!" << std::endl;
			} else if (error == 2) {
				std::cerr << "    请确认你的电脑是否安装了极域？" << std::endl;
			}

			return false;
		}

		std::cout << "[+] 成功获取驱动程序句柄。" << std::endl;
		return true;
	}

	void CloseDriverHandle() {
		if (hDevice != INVALID_HANDLE_VALUE) {
			CloseHandle (hDevice);
			hDevice = INVALID_HANDLE_VALUE;
		}
	}

	bool SendIoctl (DWORD ioctlCode, const CONFIG_PAYLOAD& payload) {
		if (!OpenDriverHandle()) {
			return false;
		}

		DWORD bytesReturned = 0;
		BOOL success = DeviceIoControl (
		                   hDevice,
		                   ioctlCode,
		                   (LPVOID)&payload,          // 输入缓冲区
		                   sizeof (CONFIG_PAYLOAD),   // 输入缓冲区大小 (0x98 字节)
		                   nullptr,                   // 输出缓冲区 (不需要)
		                   0,                         // 输出缓冲区大小
		                   &bytesReturned,
		                   nullptr                    // LPOVERLAPPED
		               );

		if (success) {
			std::cout << "[+] 命令执行成功。" << std::endl;
		} else {
			int code = GetLastError();
			std::cerr << "[X] 命令失败！错误码：" << code << std::endl;
			if (code == 2) {
				std::cerr << "  本错误可能的原因是内核程序未启动，可以忽略。";
			}
		}

		return success != 0;
	}

public:
	~TDNetFilterController() {
		CloseDriverHandle();
	}

	void Run() {
		CONFIG_PAYLOAD payload = {0};
		std::fill (payload.szListName, payload.szListName + MAX_LIST_NAME_SIZE, 0);
		strcpy (payload.szListName, "User IOCTL Config");

		int choice;
		std::cout << "==========================================\n";
		std::cout << "     极域网络控制台 By Piaoztsdy 2025\n";
		std::cout << "==========================================\n";
		std::cout << "\n注意：经过测试当教师更改网络控制时需要重新使用！\nTip: 如果选择黑名单并且不输入 IP，则等于无限制。\n\n本程序理论只对 2016(6.0) 版极域起效！\n\n";
		std::cout << "请选择操作模式:\n";
		std::cout << "  1. 白名单模式\n";
		std::cout << "  2. 黑名单模式\n";
		std::cout << "  3. 关闭过滤/重置 (具体功能暂未逆向，实测可以暂时关闭)\n";
		std::cout << "  Other. 退出\n";
		std::cout << "  > ";
		std::cin >> choice;

		if (choice == 1 || choice == 2) {
			GetIpListFromUser (&payload);
			DWORD ioctl = (choice == 1) ? IOCTL_SET_WHITELIST : IOCTL_SET_BLACKLIST;
			SendIoctl (ioctl, payload);
		} else if (choice == 3) {
			// 重置模式：清空 IP 列表，发送 IOCTL 0x120018 (Mode 0)
			payload.dwIpCount = 0;
			std::cout << "[i] 正在清空 IP 列表并发送重置命令..." << std::endl;
			SendIoctl (IOCTL_SET_RESET, payload);
		} else {
			std::cerr << "[-] 程序退出。" << std::endl;
		}
	}

	void GetIpListFromUser (CONFIG_PAYLOAD* payload) {
		std::string input;
		std::cout << "\n请输入 IP 地址列表 (使用逗号分隔，例如: 1.1.1.1,8.8.8.8):\n> ";
		// 确保清除 cin 缓冲区
		std::cin.ignore (std::numeric_limits < std::streamsize >::max(), '\n');
		std::getline (std::cin, input);

		payload->dwIpCount = 0;
		std::string token;
		std::stringstream ss (input);

		while (std::getline (ss, token, ',') && payload->dwIpCount < MAX_IP_COUNT) {
			// 清理空格
			token.erase (std::remove (token.begin(), token.end(), ' '), token.end());

			if (!token.empty()) {
				IN_ADDR ipAddr;

				// 使用inet_pton将点分十进制IP转换为网络字节序
				if (InetPton (AF_INET, token.c_str(), &ipAddr) == 1) {
					payload->IpArray[payload->dwIpCount] = ipAddr.S_un.S_addr;
					payload->dwIpCount++;
				} else {
					std::cerr << "警告: 忽略无效 IP 地址: " << token << std::endl;
				}
			}
		}

		std::cout << "[i] 已处理 " << payload->dwIpCount << " 个有效 IP 地址。" << std::endl;
	}
};

int main() {
	// 检查程序是否以管理员权限运行
	if (!IsUserAnAdmin()) {
		std::cerr << "*****************************************************" << std::endl;
		std::cerr << "***      错误：本程序必须以管理员权限运行!        ***" << std::endl;
		std::cerr << "*****************************************************" << std::endl;
		// 等待用户查看错误
		std::cout << "按任意键退出...";
		system ("pause>nul");
		return 1;
	}

	// 初始化 Winsock 用于 InetPton
	WSADATA wsaData;

	if (WSAStartup (MAKEWORD (2, 2), &wsaData) != 0) {
		std::cerr << "[X] WSAStartup 失败." << std::endl;
		return 1;
	}

	TDNetFilterController tnff;
	tnff.Run();

	WSACleanup();

	std::cout << "\n操作完成。按任意键退出...";
	system ("pause>nul");
	return 0;
}
